from io import BytesIO

import pandas as pd
from django.utils.encoding import escape_uri_path
from rest_framework.decorators import action

from backend.utils.renders import FileRender
from backend.utils.response import ApiResponse


class ExportModelMixin:
    # 导出的表头,
    excel_headers = []
    # 导出的文件名
    file_name = '未命名.xlsx'

    @action(methods=['POST'], detail=False, url_path='export', url_name='export', renderer_classes=[FileRender, ])
    def export(self, request):
        id_list = request.data.get('id_list', [])
        queryset = self.get_queryset().filter(id__in=id_list)
        data = self.get_serializer(queryset, many=True).data
        data = sorted(data, key=lambda department: id_list.index(department.get('id')))

        def remove_id(department):
            department.pop('id')
            return department

        data = map(lambda department: remove_id(department), data)
        df = pd.DataFrame(list(data))
        # 表头
        bio = BytesIO()
        writer = pd.ExcelWriter(bio, engine='xlsxwriter')
        df.to_excel(writer, header=self.excel_headers, index=False)
        writer.save()
        workbook = bio.getvalue()
        headers = {
            'Content-Type': 'application/octet-stream',
            'Content-Disposition': f'attachment;filename={escape_uri_path(self.file_name)}',
            'Access-Control-Expose-Headers': 'Content-Disposition',
        }
        return ApiResponse(data=workbook, headers=headers)
